##########################################################################################

library(dplyr)
library(ggplot2)
library(data.table)
library(RColorBrewer)
library(optparse)
library(ggpubr)
library(maftools)
library(copynumber)

##########################################################################################

option_list <- list(
    make_option(c("--sample_file"), type = "character") ,
    make_option(c("--seg_file"), type = "character") ,
    make_option(c("--sample_public_file"), type = "character") ,
    make_option(c("--seg_public_file"), type = "character") ,
    make_option(c("--class_order_file"), type = "character") ,
    make_option(c("--images_path"), type = "character")
)

if(1!=1){
    
    work_dir <- "~/20220915_gastric_multiple/dna_combinePublic/"
    sample_file <- paste(work_dir,"/baseTable/STAD_Info.addBurden.MSI_MSS.addCNVType.tsv",sep="")
    sample_public_file <- paste(work_dir,"/public_ref/combine/MutationInfo.combine.addMolecularSubType.tsv",sep="")
    seg_file <- paste(work_dir,"/titan/Titan_all_seg.final.tsv",sep="")
    seg_public_file <- paste(work_dir,"/seg_public/TCGA_use.seg",sep="")
    class_order_file <- paste(work_dir,"/config/Class_order.list",sep="")
    ######
    images_path <- paste(work_dir,"/images/cnv_burden",sep="")

}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

sample_file <- opt$sample_file
seg_file <- opt$seg_file
sample_public_file <- opt$sample_public_file
seg_public_file <- opt$seg_public_file
class_order_file <- opt$class_order_file
images_path <- opt$images_path

dir.create(images_path , recursive = T)

###########################################################################################

info <- data.frame(fread(sample_file))
seg <- data.frame(fread(seg_file))
class_order <- data.frame(fread(class_order_file))

info_public <- data.frame(fread(sample_public_file))
seg_public <- data.frame(fread(seg_public_file))

###########################################################################################
col <- c(
  brewer.pal(9,"YlGnBu")[6],
  rgb(234,106,79,alpha=255,maxColorValue=255),
  rgb(203,24,30,alpha=255,maxColorValue=255),
  rgb(255,0,0,alpha=255,maxColorValue=255)
  )

names(col) <- c("IM" , "IGC" , "DGC" , "GC")
col <- col[1:4]

###########################################################################################

dat <- seg
dat$Sample <- gsub( "S_" , "S" , dat$Sample )
info$Tumor <- gsub( "S_" , "S" , info$Tumor )
info$Normal <- gsub( "S_" , "S" , info$Normal )

dat$seg.mean <- dat$Median_logR
dat$loc.start <- dat$Start
dat$loc.end <- dat$End
dat$Tumor <- sapply( strsplit( dat$Sample , "_" ) , "[" , 1)
dat$Normal <- sapply( strsplit( dat$Sample , "_" ) , "[" , 2)

dat <- merge( dat , info[,c("Tumor" , "Class" , "Type" , "TCGA_Class")] , by = "Tumor" )

###########################################################################################

dat_public <- seg_public
dat_public <- merge( dat_public , info_public[,c("Tumor" , "Class" , "From" , "Molecular.subtype")] , by.x = "Sample" , by.y = "Tumor" )
dat_public$Normal <- dat_public$Sample

###########################################################################################
## 整合数据
# "Sample"       "Chromosome"   "Start"        "End"          "Num_Probes"  
#  "Segment_Mean" "Class" 
use_col <- c("Tumor" , "Chromosome" , "Start" , "End" , "Length.snp." , "Median_logR" , "Class" , "TCGA_Class" , "Normal" )
dat_im <- subset(dat , Class=="IM")[,use_col]
dat_gc <- subset(dat , Class!="IM" & Type != "IM + IGC + DGC" )[,use_col]
dat_nmu <- rbind(dat_im , dat_gc)
dat_nmu$From <- "NJMU"
dat_nmu <- dat_nmu[,c("Tumor" , "Chromosome" , "Start" , "End" , "Length.snp." , "Median_logR" , "Class" , "From" , "TCGA_Class" , "Normal" )]
colnames(dat_nmu) <- colnames(dat_public)

###########################################################################################

dat_combine <- rbind(dat_nmu , dat_public)
dat_combine$seg.mean <- dat_combine$Segment_Mean
dat_combine$loc.start <- dat_combine$Start
dat_combine$loc.end <- dat_combine$End
dat_combine$Tumor <- dat_combine$Sample
dat_combine <- subset( dat_combine , Chromosome %in% c(1:22) )

###########################################################################################
## 拷贝数分布图
## 定义阈值，分割扩增或缺失
#### Function
genomeFreq <- function(data,thres.gain,thres.loss,pos.unit,layout,Type, LOH_use=data.frame(0) ){

  nr <- layout[1]
  nc <- layout[2]
  cr <- nr*nc

  nProbe <- nrow(data)
  nT <- length(thres.gain)
  
  #Get plot parameters:
  op <- copynumber:::getFreqPlotParameters(type="genome",nc=nc,nr=nr,thres.gain=thres.gain,thres.loss=thres.loss)

  #Set global xlimits if not specified by user:
  if(is.null(op$xlim)){
    op$xlim <- copynumber:::getGlobal.xlim(op=op,pos.unit=pos.unit,chrom=unique(data[,1]))
  }
  
  #Adjust positions to be plotted along xaxis; i.e. get global positions, scale according to plot.unit, and get left and right pos for freq.rectangles
  #to be plotted (either continuous or 1 probe long):
  x <- copynumber:::adjustPos(position=data[,2],chromosomes=data[,1],pos.unit=pos.unit,type="genome",op=op)
  xleft <- x$xleft
  xright <- x$xright


  #Initialize:
  row=1
  clm=1
  new = FALSE

  #Division of plotting window:
  frames <- copynumber:::framedim(nr,nc)

  #One plot for each value in thres.gain/thres.loss:
  for(t in 1:nT){

    #Frame dimensions for plot t:
    fig.t <- c(frames$left[clm],frames$right[clm],frames$bot[row],frames$top[row])
    #par(fig=fig.t,new=new,oma=c(0,0,1,0),mar=op$mar)
    
    #Calculate the percentage of samples that have estimated copy number larger than thres at given position:
    freq.amp <- rowMeans(data[,-c(1:2),drop=FALSE] > thres.gain[t])*100
    freq.del <- rowMeans(data[,-c(1:2),drop=FALSE] < thres.loss[t])*100

    #Find default ylimits and at.y (tickmarks):
    op <- copynumber:::updateFreqParameters(freq.del,freq.amp,op)
    #op$assembly <- op$assembly[!op$assembly$V1 %in% c("chrX" , "chrY"),]

    #Empty plot with correct limits
    plot(1,1,type="n",ylim=c(-100,100),xlim=op$xlim,xaxs="i",main="",frame.plot=FALSE,yaxt="n",xaxt="n",ylab="",xlab="" , cex.axis = 8)

    #Add shifting white/grey pattern to backgroud to separate chromosomes:
    copynumber:::chromPattern(pos.unit,op)
    
    #main title for this plot
    title(main=Type,line=op$main.line,cex.main= 2.3,adj=0)

    #Add axes, labels and percentage lines:
    op$at.y <- seq(0,100,25)
    copynumber:::addToFreqPlot(op,type="genome")

    #Plot frequencies:
    rect(xleft=xleft,ybottom=0,xright=xright,ytop=freq.amp,col=op$col.gain,border=op$col.gain)
    rect(xleft=xleft,ybottom=0,xright=xright,ytop=-freq.del,col=op$col.loss,border=op$col.loss)

    #Add line at y=0:
    abline(h=0,lty=1,col="grey82",lwd=1.5)

    #Add line at y=0:
    abline(h=0,lty=1,col="grey82",lwd=1.5)

    #Separate chromosomes by vertical lines:
    op$chrom.lty = 1
    copynumber:::addChromlines(data[,1],xaxis="pos",unit=pos.unit,cex=op$cex.chrom,op=op)
    #copynumber:::addArmlines(data[,1],xaxis="pos",unit=pos.unit,cex=op$cex.chrom,op=op)


    #######################################
     ## LOH
    if(dim(LOH_use)[1]>1){
        LOH_use <- LOH_use[order(LOH_use$chromosome_name),]
        LOH_use$chromosome_name <- as.numeric(LOH_use$chromosome_name)
        x1 <- copynumber:::adjustPos(position=LOH_use[,3],chromosomes=LOH_use[,2],pos.unit=pos.unit,type="genome",op=op)
        xleft_loh <- (x1$xleft + x1$xright)/2
        x1 <- copynumber:::adjustPos(position=LOH_use[,4],chromosomes=LOH_use[,2],pos.unit=pos.unit,type="genome",op=op)
        xright_loh <- (x1$xleft + x1$xright)/2

        line <- rep(1.2,dim(LOH_use)[1])
        line[which(LOH_use$hgnc_symbol=="TBL1XR1")] <- 2.2
        line[which(LOH_use$hgnc_symbol=="RNF43")] <- 0.2
        line[which(LOH_use$hgnc_symbol=="CDH1")] <- 2.2
        line[which(LOH_use$hgnc_symbol=="ERBB2")] <- 0.2

        ## Plot LOH:
        loh_colr <- brewer.pal(9,"Set1")[5]
        rect(xleft=xleft_loh,ybottom=par("usr")[3] ,xright=xright_loh,ytop=(par("usr")[4]),border=loh_colr)
        mtext(LOH_use[,1], side = 3,  at = (xleft_loh + xright_loh)/2  , line = line, cex= 0.6 ,col=loh_colr)

    }

  }#endfor

}#end function

###########################################################################################
## IGC中，CIN和CS的区域
## DGC中，CIN和GS的区域

Type_all <- c("IGC","DGC")
cnv_Type <- c("CIN","GS")

pos.unit = "bp"
thres.gain=0.2
thres.loss=-0.2

for(laruen in Type_all){

  images_name <- paste0(images_path,"/CopyRatioDis.",laruen,".NJMU.pdf")
  pdf(images_name , height = 5 , width = 10)

  ## 三张图放一起
  par(mfrow=c(2,1),mar=c(2,2,3,1))

  for(cnv_t in cnv_Type){
    print(cnv_t)

    ## 读文件
    seg <- subset( dat_combine , Molecular.subtype == cnv_t & Class == laruen )
    seg <- seg[,c("Sample" , "Chromosome" , "Start" , "End" , "Num_Probes" , "Segment_Mean")]
    colnames(seg) <- c("sampleID" , "chrom" , "start.pos" , "end.pos" , "n.probes" , "mean")
    seg <- subset( seg , !(chrom %in% c("X" , "Y")) )

    ## 得到的seg的频率分布
    ## 分染色体，将所有的样本的start和end都收集起来
    ## 看每个样本在对应的pos处的logRatio
    data = seg
    data = copynumber:::getFreqData(data)
    data[, 1] <- as.numeric(data[, 1])
    ## NJMU的合并一个人的多个文件
    tmp_njmu <- data.frame(chr = data$chr , pos = data$pos)
    for(id in unique(info$Patient)){
      tumors <- subset( info , Patient == id & Class == laruen & Type != "IM + IGC + DGC" & TCGA_Class == cnv_t  )$Tumor
      if(length(tumors) > 1){
        tmp <- data.frame( value = apply(data[,tumors] , 1 , median) )
        colnames(tmp) <- id
        tmp_njmu <- cbind(tmp_njmu , tmp)
      }else if(length(tumors) == 1){
        tmp <- data.frame( value = data[,tumors])
        colnames(tmp) <- id
        tmp_njmu <- cbind(tmp_njmu , tmp)
      }
    }

    data_use <- tmp_njmu
    tmp_tcga <- data[,grep( "TCGA" , colnames(data))]
    data_use <- cbind(tmp_njmu , tmp_tcga)

    chrom <- copynumber:::checkChrom(data = data_use, segments = NULL, chrom = NULL)
    type <- ifelse(is.null(chrom), "genome", "bychrom")
    ## 画图
    genomeFreq(data = data_use,thres.gain=thres.gain , thres.loss = thres.loss ,pos.unit="bp",layout = c(1,1),Type = cnv_t)
  }

  dev.off()
}


for(laruen in Type_all){

  images_name <- paste0(images_path,"/CopyRatioDis.",laruen,".TCGA.pdf")
  pdf(images_name , height = 5 , width = 10)

  ## 三张图放一起
  par(mfrow=c(2,1),mar=c(2,2,3,1))

  for(cnv_t in cnv_Type){
    print(cnv_t)

    ## 读文件
    seg <- subset( dat_combine , Molecular.subtype == cnv_t & Class == laruen )
    seg <- seg[,c("Sample" , "Chromosome" , "Start" , "End" , "Num_Probes" , "Segment_Mean")]
    colnames(seg) <- c("sampleID" , "chrom" , "start.pos" , "end.pos" , "n.probes" , "mean")
    seg <- subset( seg , !(chrom %in% c("X" , "Y")) )

    ## 得到的seg的频率分布
    ## 分染色体，将所有的样本的start和end都收集起来
    ## 看每个样本在对应的pos处的logRatio
    data = seg
    data = copynumber:::getFreqData(data)
    data[, 1] <- as.numeric(data[, 1])

    tmp_njmu <- data.frame(chr = data$chr , pos = data$pos)
    tmp_tcga <- data[,grep( "TCGA" , colnames(data))]
    data_use <- cbind(tmp_njmu , tmp_tcga)

    chrom <- copynumber:::checkChrom(data = data_use, segments = NULL, chrom = NULL)
    type <- ifelse(is.null(chrom), "genome", "bychrom")
    ## 画图
    genomeFreq(data = data_use,thres.gain=thres.gain , thres.loss = thres.loss ,pos.unit="bp",layout = c(1,1),Type = cnv_t)
  }

  dev.off()

}

###########################################################################################

Type_all <- c("IM","IGC","DGC")
pos.unit = "bp"
thres.gain=0.2
thres.loss=-0.2

images_name <- paste0(images_path,"/CopyRatioDis.pdf")
pdf(images_name , height = 7/1.5 , width = 11.2/1.5)

## 三张图放一起
par(mfrow=c(3,1),mar=c(2,2,3,1))
for(Type in Type_all){
  print(Type)
  class <- Type

  ## 读文件
  seg <- subset( dat_combine , Class == Type )
  seg <- seg[,c("Sample" , "Chromosome" , "Start" , "End" , "Num_Probes" , "Segment_Mean")]
  colnames(seg) <- c("sampleID" , "chrom" , "start.pos" , "end.pos" , "n.probes" , "mean")
  seg <- subset( seg , !(chrom %in% c("X" , "Y")) )

  ## 得到的seg的频率分布
  ## 分染色体，将所有的样本的start和end都收集起来
  ## 看每个样本在对应的pos处的logRatio
  data = seg
  data = copynumber:::getFreqData(data)
  data[, 1] <- as.numeric(data[, 1])

  if(Type!="IM"){
    tmp_njmu <- data.frame(chr = data$chr , pos = data$pos)
    tmp_tcga <- data[,grep( "TCGA" , colnames(data))]
    data_use <- cbind(tmp_njmu , tmp_tcga)
  }else{
    data_use <- data
  }

  chrom <- copynumber:::checkChrom(data = data_use, segments = NULL, chrom = NULL)
  type <- ifelse(is.null(chrom), "genome", "bychrom")

  ## 画图
  genomeFreq(data = data_use,thres.gain=thres.gain , thres.loss = thres.loss ,pos.unit="bp",layout = c(1,1),Type = Type)
}

dev.off()
